"""
22_hnsw_indexing.py
===================
HNSW (Hierarchical Navigable Small World) Indexing in Neo4j

This module covers:
1. Understanding HNSW algorithm
2. Creating vector indexes
3. Index configuration and tuning
4. Similarity search operations
5. Performance optimization
6. How Neo4j stores vectors internally

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import time
import math
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class SimilarityFunction(Enum):
    """Supported similarity functions."""
    COSINE = "cosine"
    EUCLIDEAN = "euclidean"


@dataclass
class VectorIndexConfig:
    """Configuration for vector index."""
    name: str
    label: str
    property_name: str
    dimensions: int
    similarity_function: SimilarityFunction = SimilarityFunction.COSINE
    m: int = 16  # Max connections per layer
    ef_construction: int = 64  # Size of dynamic candidate list


class HNSWIndexManager:
    """
    Manages HNSW vector indexes in Neo4j.
    
    HNSW (Hierarchical Navigable Small World) is a graph-based 
    approximate nearest neighbor search algorithm that provides:
    
    - Sub-linear search time: O(log n) average
    - High recall rates: >95% typical
    - Efficient index construction
    - Good scalability to millions of vectors
    
    How HNSW Works:
    ===============
    
    1. HIERARCHICAL LAYERS
       - Multiple layers of decreasing density
       - Top layer: Few nodes, long-range connections
       - Bottom layer: All nodes, local connections
       
       Layer 3:  [A]------------------[B]
                   \\                  /
       Layer 2:  [A]---[C]-------[B]
                  |     |         |
       Layer 1:  [A]-[D]-[C]-[E]-[B]
                  | | | | | | | | |
       Layer 0:  [A][D][C][F][E][G][B] (all nodes)
    
    2. SEARCH PROCESS
       - Start at top layer
       - Greedily navigate to nearest neighbor
       - Drop to lower layer and repeat
       - Final layer gives precise results
       
    3. KEY PARAMETERS
       - M: Max connections per node (tradeoff: memory vs accuracy)
       - ef_construction: Build-time search width (higher = better index)
       - ef_search: Query-time search width (higher = better recall)
    
    Internal Storage:
    ================
    Neo4j stores HNSW indexes in a specialized format:
    
    - Vectors stored as native arrays
    - Layer structure in dedicated index files
    - Connections encoded as adjacency lists
    - Separate files for each layer level
    
    Performance Characteristics:
    ===========================
    - Index build: O(n log n)
    - Single query: O(log n)
    - Memory: O(n * M * layers)
    - Disk: Similar to memory (persistent index)
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
    
    # =========================================================
    # INDEX CREATION
    # =========================================================
    
    def create_vector_index(
        self,
        config: VectorIndexConfig
    ) -> Dict[str, Any]:
        """
        Create an HNSW vector index.
        
        Args:
            config: Index configuration
            
        Returns:
            Creation result
            
        Example:
            config = VectorIndexConfig(
                name="chunk_embeddings",
                label="Chunk",
                property_name="embedding",
                dimensions=384,  # for all-MiniLM-L6-v2
                similarity_function=SimilarityFunction.COSINE
            )
        """
        query = f"""
        CREATE VECTOR INDEX {config.name} IF NOT EXISTS
        FOR (n:{config.label})
        ON (n.{config.property_name})
        OPTIONS {{
            indexConfig: {{
                `vector.dimensions`: {config.dimensions},
                `vector.similarity_function`: '{config.similarity_function.value}'
            }}
        }}
        """
        
        with self.driver.session() as session:
            try:
                session.run(query)
                
                # Wait for index to come online
                return self._wait_for_index(config.name)
            except Exception as e:
                return {"status": "error", "message": str(e)}
    
    def _wait_for_index(
        self,
        index_name: str,
        timeout: int = 60
    ) -> Dict[str, Any]:
        """
        Wait for index to become online.
        
        Args:
            index_name: Name of index
            timeout: Maximum wait time in seconds
            
        Returns:
            Index status
        """
        query = """
        SHOW INDEXES
        WHERE name = $name
        YIELD name, state, populationPercent
        RETURN name, state, populationPercent
        """
        
        start = time.time()
        while time.time() - start < timeout:
            with self.driver.session() as session:
                result = session.run(query, name=index_name)
                record = result.single()
                
                if record:
                    if record["state"] == "ONLINE":
                        return {
                            "status": "online",
                            "name": index_name,
                            "population": record["populationPercent"]
                        }
                    elif record["state"] == "FAILED":
                        return {
                            "status": "failed",
                            "name": index_name
                        }
            
            time.sleep(1)
        
        return {"status": "timeout", "name": index_name}
    
    # =========================================================
    # INDEX MANAGEMENT
    # =========================================================
    
    def list_vector_indexes(self) -> List[Dict[str, Any]]:
        """
        List all vector indexes.
        
        Returns:
            List of vector index information
        """
        query = """
        SHOW INDEXES
        WHERE type = 'VECTOR'
        YIELD name, labelsOrTypes, properties, state, 
              populationPercent, options
        RETURN name, labelsOrTypes, properties, state,
               populationPercent, options
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [record.data() for record in result]
    
    def get_index_stats(self, index_name: str) -> Dict[str, Any]:
        """
        Get detailed statistics for a vector index.
        
        Args:
            index_name: Name of the index
            
        Returns:
            Index statistics
        """
        query = """
        SHOW INDEXES
        WHERE name = $name
        YIELD name, type, labelsOrTypes, properties,
              state, populationPercent, options, 
              indexProvider, owningConstraint
        RETURN *
        """
        
        with self.driver.session() as session:
            result = session.run(query, name=index_name)
            record = result.single()
            
            if record:
                data = record.data()
                
                # Extract vector-specific options
                options = data.get('options', {})
                index_config = options.get('indexConfig', {})
                
                return {
                    "name": data['name'],
                    "type": data['type'],
                    "labels": data['labelsOrTypes'],
                    "property": data['properties'],
                    "state": data['state'],
                    "population_percent": data['populationPercent'],
                    "dimensions": index_config.get('vector.dimensions'),
                    "similarity_function": index_config.get('vector.similarity_function'),
                    "provider": data['indexProvider']
                }
            
            return None
    
    def drop_index(self, index_name: str) -> bool:
        """
        Drop a vector index.
        
        Args:
            index_name: Name of index to drop
            
        Returns:
            True if dropped
        """
        query = f"DROP INDEX {index_name} IF EXISTS"
        
        with self.driver.session() as session:
            try:
                session.run(query)
                return True
            except Exception as e:
                print(f"   Error dropping index: {e}")
                return False
    
    # =========================================================
    # SIMILARITY SEARCH
    # =========================================================
    
    def vector_similarity_search(
        self,
        index_name: str,
        query_vector: List[float],
        top_k: int = 10,
        min_score: float = 0.0
    ) -> List[Dict[str, Any]]:
        """
        Perform vector similarity search using HNSW index.
        
        This is the core operation for semantic search in RAG.
        
        Args:
            index_name: Name of vector index
            query_vector: Query embedding vector
            top_k: Number of results to return
            min_score: Minimum similarity score filter
            
        Returns:
            List of similar nodes with scores
            
        How it works internally:
        1. Query vector enters at top HNSW layer
        2. Greedy search finds approximate nearest neighbor
        3. Algorithm descends to next layer
        4. Repeat until bottom layer
        5. Expand search at bottom for final candidates
        6. Return top-k by similarity score
        """
        query = """
        CALL db.index.vector.queryNodes($index_name, $k, $vector)
        YIELD node, score
        WHERE score >= $min_score
        RETURN 
            node.id AS id,
            node.content AS content,
            score,
            labels(node) AS labels
        ORDER BY score DESC
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                k=top_k,
                vector=query_vector,
                min_score=min_score
            )
            return [record.data() for record in result]
    
    def hybrid_search(
        self,
        index_name: str,
        query_vector: List[float],
        keyword_filter: str,
        top_k: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Combine vector similarity with keyword filtering.
        
        This provides better precision by combining semantic
        and lexical matching.
        
        Args:
            index_name: Vector index name
            query_vector: Query embedding
            keyword_filter: Keyword to filter by
            top_k: Results to return
            
        Returns:
            Filtered similar nodes
        """
        query = """
        CALL db.index.vector.queryNodes($index_name, $k * 3, $vector)
        YIELD node, score
        WHERE toLower(node.content) CONTAINS toLower($keyword)
        WITH node, score
        ORDER BY score DESC
        LIMIT $k
        RETURN 
            node.id AS id,
            node.content AS content,
            score
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                k=top_k,
                vector=query_vector,
                keyword=keyword_filter
            )
            return [record.data() for record in result]
    
    def search_with_graph_context(
        self,
        index_name: str,
        query_vector: List[float],
        top_k: int = 5,
        context_hops: int = 1
    ) -> List[Dict[str, Any]]:
        """
        Vector search with graph neighborhood expansion.
        
        After finding semantically similar nodes, expands to
        include connected nodes for richer context.
        
        Args:
            index_name: Vector index name
            query_vector: Query embedding
            top_k: Initial results
            context_hops: Relationship hops to expand
            
        Returns:
            Results with graph context
        """
        query = f"""
        CALL db.index.vector.queryNodes($index_name, $k, $vector)
        YIELD node, score
        WITH node, score
        ORDER BY score DESC
        LIMIT $k
        
        // Expand to get connected context
        OPTIONAL MATCH (node)-[r*1..{context_hops}]-(related)
        WHERE NOT related:Chunk  // Don't include other chunks
        
        WITH node, score, 
             collect(DISTINCT {{
                 label: labels(related)[0],
                 name: related.name,
                 relationship: type(r[0])
             }}) AS context
        
        RETURN 
            node.id AS id,
            node.content AS content,
            score,
            context
        """
        
        with self.driver.session() as session:
            result = session.run(
                query,
                index_name=index_name,
                k=top_k,
                vector=query_vector
            )
            return [record.data() for record in result]
    
    # =========================================================
    # PERFORMANCE BENCHMARKING
    # =========================================================
    
    def benchmark_search(
        self,
        index_name: str,
        query_vectors: List[List[float]],
        top_k: int = 10
    ) -> Dict[str, Any]:
        """
        Benchmark search performance.
        
        Args:
            index_name: Index to benchmark
            query_vectors: List of query vectors
            top_k: Results per query
            
        Returns:
            Performance statistics
        """
        times = []
        
        for vector in query_vectors:
            start = time.perf_counter()
            self.vector_similarity_search(index_name, vector, top_k)
            elapsed = (time.perf_counter() - start) * 1000
            times.append(elapsed)
        
        return {
            "queries": len(query_vectors),
            "top_k": top_k,
            "min_ms": min(times),
            "max_ms": max(times),
            "avg_ms": sum(times) / len(times),
            "p50_ms": sorted(times)[len(times) // 2],
            "p95_ms": sorted(times)[int(len(times) * 0.95)],
            "p99_ms": sorted(times)[int(len(times) * 0.99)] if len(times) > 100 else max(times)
        }
    
    # =========================================================
    # UTILITY METHODS
    # =========================================================
    
    def count_indexed_vectors(self, label: str, property_name: str) -> int:
        """
        Count nodes with vector embeddings.
        
        Args:
            label: Node label
            property_name: Embedding property name
            
        Returns:
            Count of indexed vectors
        """
        query = f"""
        MATCH (n:{label})
        WHERE n.{property_name} IS NOT NULL
        RETURN count(n) as count
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return result.single()["count"]
    
    def sample_vectors(
        self,
        label: str,
        property_name: str,
        sample_size: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Get sample vectors from the database.
        
        Args:
            label: Node label
            property_name: Embedding property
            sample_size: Number to sample
            
        Returns:
            Sample vectors with metadata
        """
        query = f"""
        MATCH (n:{label})
        WHERE n.{property_name} IS NOT NULL
        RETURN 
            n.id AS id,
            n.{property_name} AS vector,
            size(n.{property_name}) AS dimensions
        LIMIT $limit
        """
        
        with self.driver.session() as session:
            result = session.run(query, limit=sample_size)
            return [record.data() for record in result]


# ============================================================
# HNSW ALGORITHM EXPLANATION
# ============================================================

def explain_hnsw_algorithm():
    """
    Detailed explanation of HNSW algorithm.
    
    This function provides educational content about how
    HNSW works internally.
    """
    explanation = """
    ╔══════════════════════════════════════════════════════════════════╗
    ║          HNSW: Hierarchical Navigable Small World                ║
    ╠══════════════════════════════════════════════════════════════════╣
    
    WHAT IS HNSW?
    ─────────────
    HNSW is an algorithm for Approximate Nearest Neighbor (ANN) search.
    It builds a multi-layer graph structure that enables fast similarity
    searches in high-dimensional spaces.
    
    KEY INSIGHT: Instead of comparing query to ALL vectors (O(n)),
    HNSW uses graph navigation to find similar vectors in O(log n).
    
    
    STRUCTURE: HIERARCHICAL LAYERS
    ──────────────────────────────
    
    Layer 2 (sparse):    [A]━━━━━━━━━━━━━━━━━━━[B]
                          │                     │
    Layer 1 (medium):    [A]━━━━[C]━━━━━━[D]━━━[B]
                          │     │         │     │
    Layer 0 (dense):     [A]-[E]-[C]-[F]-[D]-[G]-[B]
                          All nodes at this level
    
    • Top layers: Few nodes, LONG connections (fast coarse search)
    • Bottom layers: All nodes, LOCAL connections (precise search)
    • Each node exists at a random level (exponentially distributed)
    
    
    SEARCH ALGORITHM
    ────────────────
    
    1. START at top layer with entry point
    2. GREEDY SEARCH: Move to neighbor closest to query
    3. Repeat until no closer neighbor found
    4. DROP DOWN to next layer
    5. Repeat steps 2-4 until layer 0
    6. EXPAND SEARCH at layer 0 for final results
    
    Example search for query Q:
    
    Layer 2:  [A]══════════════[B]  (Start at A, B is closer to Q)
                   ↓
    Layer 1:  [A]──[C]──[D]──[B]    (From B, D is closer)
                        ↓
    Layer 0:  ...[F]─[D]─[G]...     (From D, expand to find F,G)
    
    Result: Return D, F, G as nearest neighbors
    
    
    KEY PARAMETERS
    ──────────────
    
    M (Max Connections):
    • Controls graph density
    • Higher M = more memory, better recall
    • Typical values: 16-64
    • Trade-off: Memory vs Search Quality
    
    ef_construction (Build-time ef):
    • Search width during index construction
    • Higher = better index quality, slower build
    • Typical values: 64-200
    • Only affects build time, not search time
    
    ef_search (Query-time ef):
    • Search width during queries
    • Higher = better recall, slower queries
    • Can be tuned per-query
    • Typical values: 50-200
    
    
    PERFORMANCE CHARACTERISTICS
    ──────────────────────────
    
    Time Complexity:
    • Build: O(n × log(n))
    • Search: O(log(n))
    
    Space Complexity:
    • O(n × M × average_layers)
    • Typical: ~1.3 × raw vector storage
    
    Recall vs Speed:
    ┌─────────────┬─────────────┬─────────────┐
    │ ef_search   │ Recall      │ Query Time  │
    ├─────────────┼─────────────┼─────────────┤
    │ 10          │ ~85%        │ 0.5ms       │
    │ 50          │ ~95%        │ 2ms         │
    │ 100         │ ~98%        │ 5ms         │
    │ 200         │ ~99%        │ 10ms        │
    └─────────────┴─────────────┴─────────────┘
    
    
    WHY HNSW FOR NEO4J RAG?
    ──────────────────────
    
    1. SPEED: Sub-millisecond searches on millions of vectors
    2. ACCURACY: High recall rates (>95% typical)
    3. SCALABILITY: Handles growing knowledge bases
    4. INTEGRATION: Combines with graph traversal
    5. PERSISTENCE: Index survives database restarts
    
    ╚══════════════════════════════════════════════════════════════════╝
    """
    print(explanation)


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("HNSW Vector Indexing - Demonstration")
    print("=" * 60)
    
    # Explain algorithm
    explain_hnsw_algorithm()
    
    # Practical demo
    print("\n📌 Practical HNSW Index Operations")
    print("-" * 40)
    
    try:
        driver = get_driver()
        manager = HNSWIndexManager(driver)
        
        # List existing indexes
        print("\n   Existing Vector Indexes:")
        indexes = manager.list_vector_indexes()
        if indexes:
            for idx in indexes:
                print(f"   - {idx['name']}: {idx['state']}")
        else:
            print("   No vector indexes found")
        
        # Create a test index
        print("\n   Creating test vector index...")
        config = VectorIndexConfig(
            name="test_embedding_index",
            label="TestNode",
            property_name="embedding",
            dimensions=384,
            similarity_function=SimilarityFunction.COSINE
        )
        
        result = manager.create_vector_index(config)
        print(f"   Result: {result}")
        
        # Get index stats
        stats = manager.get_index_stats("test_embedding_index")
        if stats:
            print(f"\n   Index Stats:")
            print(f"   - Dimensions: {stats['dimensions']}")
            print(f"   - Similarity: {stats['similarity_function']}")
            print(f"   - State: {stats['state']}")
        
        # Cleanup
        print("\n   Cleaning up test index...")
        manager.drop_index("test_embedding_index")
        
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ HNSW indexing demonstration complete!")
    print("=" * 60)
